#!/usr/bin/env python

import os, sys, shutil
import numpy as np
from pylab import *
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits.axes_grid import AxesGrid
from matplotlib import patches
#import matplotlib.gridspec as gridspec
import pdb
from pywolb.utilities import diversetools
from pywolb.utilities import numpytools

moviecodecs = {'mpg': 'mpeg4', \
               'mpeg': 'mpeg4', \
               'mp4': 'mpeg4', \
               'h.264': 'x264',\
               'wmv': 'wmv2'}

def draw_grid(ax,nx,ny,color='k',offset=0.):
    # vertical grid lines:
    for i in range(nx-1):
        ax.plot( [i+1-offset, i+1-offset], [0, ny], "%s-" % color, linewidth=mpl.rcParams['axes.linewidth'] )
    # horizontal grid lines:
    for i in range(ny-1):
        ax.plot( [0, nx], [i+1, i+1], "%s-" % color, linewidth=mpl.rcParams['axes.linewidth'] )

def customize_plot(ax,nPops,alleles,generation,offset=0.):
    nAlleles = len(alleles)
    draw_grid(ax,nPops,nAlleles,offset=offset)

    ax.set_xlim( 0, nPops )
    ax.set_xticks( arange(0.5, nPops+0.5, 2) )
    ax.set_xticklabels( arange(1, nPops, 2) )
    ax.set_xlabel( "Population")

    ax.set_ylim( 0, nAlleles )
    ax.set_yticks( arange(0.5, nAlleles+0.5, 1) )
    ax.set_yticklabels( alleles )
    if nAlleles > 1:
        ax.set_ylabel( "Allele" )

    ax.text(1., 1.1, "%d" % generation, horizontalalignment='right', transform=ax.transAxes) 
    

def customize_plot2D(ax,sh,generation,offset=0.):
    nX,nY = sh
    draw_grid(ax,nX,nY,offset=offset)

    ax.set_xlim( 0, nX )
    ax.set_xticks( [] )

    ax.set_ylim( 0, nY )
    ax.set_yticks( [] )
    
    ax.text(1., 1.1, "%d" % generation, horizontalalignment='right', transform=ax.transAxes)    
    

def customize_hybridzone_plot(g,sh,generation,color='k',offset=0.):
    nX,nY = sh
    axls,axhz,axrs = g[0],g[1],g[2]
    draw_grid(axhz,nX,nY,color=color,offset=offset)

    axhz.set_xlim( 0, nX )
    axhz.set_xticks( [] )
    axhz.set_xlabel( 'hybrid zone' )
    axhz.set_ylim( 0, nY )
    axhz.set_yticks( [] )
    axhz.text(1., 1.1, "%d" % generation, horizontalalignment='right', transform=axhz.transAxes)
    
    axls.set_xlim( 0,1 )
    axls.set_xticks( [] )
    axls.set_xlabel( 'left source' )
    axrs.set_xlim( 0,1 )
    axrs.set_xticks( [] )
    axrs.set_xlabel( 'right source' )


def animate(mp, alleles=['W'], outfile='animation.mpg', play=False, offset=0.):        
    try:
        global rcParams
        
        rcParams['xtick.major.size'] = 0.
        rcParams['ytick.major.size'] = 0.

        files = []
        
#        alleles = ['W', 'A1']
        nAlleles = len(alleles)
        nPops = mp.number_of_populations
        freqs = []
        for a in alleles:
            freqs.append( np.array(mp.get_hsum(a)) )
        freqs = np.array(freqs)
    #    w = np.array( mp.get_hsum('W') )    # get Wolbachia frequencies across history and populations
        generations = mp.get_generations()
        
        fig = plt.figure(figsize=(nPops+0.5,nAlleles+1.))
        fig.subplots_adjust(left=0.08, right=0.92)
        
        grid = AxesGrid(fig, 111,
                        nrows_ncols = (1,1),
                        aspect = True,
                        axes_pad = 0.,
                        share_all = False,
                        label_mode = 'L',
                        cbar_mode = 'single',
                        cbar_location = 'right',
                        cbar_pad = 0.5
                        )
        axi = grid[0]
        axcb = grid.cbar_axes[0]

        for i in range(len(generations)):
            axi.cla()
            g = generations[i]
            z = freqs[:,i,:]
            im = axi.pcolor( z, cmap='gray_r', vmin=0., vmax=1. )
            customize_plot(axi, nPops, alleles, g, offset=0.01)
            if i==0:    # draw the colorbar just once
                plt.colorbar(im, cax=axcb)
                axcb.colorbar(im)
                axcb.set_yticks( mlab.frange(0.,1.,0.25)  )
                axcb.set_yticklabels( [' 0.0', '', ' 0.5', '', ' 1.0'] )
                axcb.axis["right"].toggle(ticks=True, ticklabels=True)
    #            for s in axcb.spines:
    #                axcb.spines[s].set_linewidth(0.)

            plt.draw()
            plt.show()
            fname = '_tmp%05d.png' % g
            print 'Saving frame', fname
            savefig(fname)
            files.append(fname)

        plt.close(fig)
        print "Making movie %s - this make take a while" % outfile
        mencode = os.system("mencoder 'mf://_tmp*.png' -mf type=png:fps=10 -ovc lavc -lavcopts vcodec=wmv2 -oac copy -o %s" % outfile)    # store return value in variable mencode
        print 'mencode:', mencode
        #os.system("convert _tmp*.png animation.mng")


    finally:
        # cleanup
        if files:
            for fname in files: os.remove(fname)
        # play animation if mencoder was successful and play toggle is True
        if not mencode and play:
            os.system("mplayer -fixed-vo -idle %s" % outfile)

def animate2D(mp, sh, allele='W', outfile='animation2D.mpg', play=False, offset=0.):
    try:
        global rcParams
        
        rcParams['xtick.major.size'] = 0.
        rcParams['ytick.major.size'] = 0.

        files = []
        
        nX,nY = sh   # shape of the metapopulation
        nPops = mp.number_of_populations
        if not (nX*nY == nPops):
            raise ValueError, "invalid shape of metapopulation"
        freqs = mp.get_hsum(allele)
        generations = mp.get_generations()
        
        fig = plt.figure(figsize=(nX+1.,nY+1.))
        fig.subplots_adjust(left=0.1, right=0.8)
        
        grid = AxesGrid(fig, 111,
                        nrows_ncols = (1,1),
                        aspect = True,
                        axes_pad = 0.,
                        share_all = False,
                        label_mode = 'L',
                        cbar_mode = 'single',
                        cbar_location = 'right',
                        cbar_pad = 0.5
                        )
        axi = grid[0]
        axcb = grid.cbar_axes[0]

        for i in range(len(generations)):
            axi.cla()
            g = generations[i]
    #        z = np.array([w[i]])
    #        im = axi.imshow( z, cmap='gray_r', interpolation='nearest', vmin=0., vmax=1. )   # use reverse grayscale color map, using vmin and vmax ensures that the colorbar spans the entire grayscale range
            z = freqs[i]
            im = axi.pcolor( z.reshape((nY,nX))[::-1,:], cmap='gray_r', vmin=0., vmax=1. )  # reshape...: make sure that first population is drawn in top left corner, last population in bottom right corner
            customize_plot2D(axi, sh, g, offset=offset)
            if i==0:    # draw the colorbar just once
                plt.colorbar(im, cax=axcb)
                axcb.colorbar(im)
                axcb.set_yticks( mlab.frange(0.,1.,0.25)  )
                axcb.set_yticklabels( [' 0.0', '', ' 0.5', '', ' 1.0'] )
                axcb.axis["right"].toggle(ticks=True, ticklabels=True)
    #            for s in axcb.spines:
    #                axcb.spines[s].set_linewidth(0.)

            plt.draw()
            plt.show()
            fname = '_tmp%05d.png' % g
            print 'Saving frame', fname
            savefig(fname)
            files.append(fname)

        plt.close(fig)
        print "Making movie %s - this make take a while" % outfile
        mencode = os.system("mencoder 'mf://_tmp*.png' -mf type=png:fps=10 -ovc lavc -lavcopts vcodec=wmv2 -oac copy -o %s" % outfile)    # store return value in variable mencode
        print 'mencode:', mencode
        #os.system("convert _tmp*.png animation.mng")


    finally:
        # cleanup
        if files:
            for fname in files: os.remove(fname)
        # play animation if mencoder was successful and play toggle is True
        if not mencode and play:
            os.system("mplayer -fixed-vo -idle %s" % outfile)

def animate_hybrid_zone(mp, sh, allele='W', outfile='hybridzone.mpg', play=False, color='r', offset=0.):
    try:
        global rcParams
        
        rcParams['xtick.major.size'] = 0.
        rcParams['ytick.major.size'] = 0.

        files = []
        
        nX,nY = sh                      # shape of the hybrid zone
        nHZ = nX*nY                     # number of populations in the hybrid zone
        nT = mp.number_of_populations   # total number of populations (hybrid zone + source populations)
        if not (nX*nY == nT-2):
            raise ValueError, "invalid shape of metapopulation or total number of populations"
        freqs = mp.get_hsum(allele)
        generations = mp.get_generations()
        
        fig = plt.figure(figsize=(2*nX+2.,2*nY))
        fig.subplots_adjust(left=0.1, right=0.85)
        
        grid = AxesGrid(fig, 111,
                        nrows_ncols = (1,3),
                        aspect = True,
                        axes_pad = 0.25,
                        share_all = False,
                        label_mode = 'L',
                        cbar_mode = 'single',
                        cbar_location = 'right',
                        cbar_pad = 0.5,
                        cbar_size= "10%"
                        )
        axls = grid[0]     # left source population
        axhz = grid[1]     # hybrid zone
        axrs = grid[2]     # right source population
        axcb = grid.cbar_axes[0]

        for i in range(len(generations)):
            axls.cla()
            axhz.cla()
            axrs.cla()
            g = generations[i]
            z = freqs[i]
#            pdb.set_trace()            
            imls = axls.pcolor( z[0]*np.ones((2,2)), cmap='gray_r', vmin=0., vmax=1. )
            imhz = axhz.pcolor( z[1:nT-1].reshape((nY,nX))[::-1,:], cmap='gray_r', vmin=0., vmax=1. )  # reshape...: make sure that first population is drawn in top left corner, last population in bottom right corner
            imrs = axrs.pcolor( z[nT-1]*np.ones((2,2)), cmap='gray_r', vmin=0., vmax=1. )
            customize_hybridzone_plot(grid, sh, g, color=color, offset=offset)
            if i==0:    # draw the colorbar just once
#                plt.colorbar(imhz, cax=axcb)
                axcb.colorbar(imhz)
                axcb.set_yticks( mlab.frange(0.,1.,0.25)  )
                axcb.set_yticklabels( [' 0.0', '', ' 0.5', '', ' 1.0'] )
                axcb.axis["right"].toggle(ticks=True, ticklabels=True)
    #            for s in axcb.spines:
    #                axcb.spines[s].set_linewidth(0.)

            plt.draw()
            plt.show()
            fname = '_tmp%05d.png' % g
            print 'Saving frame', fname
            savefig(fname)
            files.append(fname)

        plt.close(fig)
        print "Making movie %s - this make take a while" % outfile
        mencode = os.system("mencoder 'mf://_tmp*.png' -mf type=png:fps=10 -ovc lavc -lavcopts vcodec=wmv2 -oac copy -o %s" % outfile)    # store return value in variable mencode
#        print 'mencode:', mencode
        #os.system("convert _tmp*.png animation.mng")

#    except:
#        # cleanup
#        if files:
#            for fname in files: os.remove(fname)
#        raise
        
    finally:
        # cleanup
        if files:
            for fname in files: os.remove(fname)
        # play animation if mencoder was successful and play toggle is True
        if not mencode and play:
            os.system("mplayer -fixed-vo -idle %s" % outfile)

def make_ticklabels_invisible(fig):
    for i, ax in enumerate(fig.axes):
        # ax.text(0.5, 0.5, "ax%d" % (i+1), va="center", ha="center")
        for tl in ax.get_xticklabels() + ax.get_yticklabels():
            tl.set_visible(False)

def popviz(ax, pop, loci='all', generation=-1, edgecolor='k', linewidth=1, withlabels=False, showaxes=False, fontsize=12, fontweight='normal', withstar=False, starloc='P1', staralpha=1.):
    hist = pop.history
    if not loci == 'all':
        for l in pop.labels['axes']:
            if l not in loci:
                hist = hist._get_axis_slice(l)
    loci = hist.labels['axes']
    labels = hist.labels['elements']
    if generation == -1:
        generation = pop.generation         # 'current generation
    d = hist.getf(generation)
    sums = numpytools.pseudoflat2(d)
    nLoci, nAlleles = np.shape(sums)
    ax.set_aspect('equal')
    ax.set_xlim(0,nLoci)
    ax.set_ylim(0,nAlleles)
    if not showaxes:
        ax.set_axis_off()
    i = 0
    for locus in sums:                   # locus
        j = 0
        for allele in locus:             # allele
            if not np.isnan(allele):
                g = 1.-allele  #round(1. - allele, 3)          # grayscale value
                ax.add_patch(patches.Rectangle((float(i),float(j)), width=1., height=1., ec=edgecolor, fc=str(g), lw=linewidth, clip_on=False))
                if withstar:
                    if labels[i][j]==starloc:
                        ax.add_patch( star(n=10, rin=0.05, rout=0.15, sh=(i+0.5,j+0.5),fc='yellow',ec='k',lw=1, alpha=staralpha) )
                if withlabels:
                    labelcolor = 'k'
                    if g < 0.25:
                        labelcolor = 'w'
                    ax.text(i+0.5,j+0.5, labels[i][j], ha='center', va='center', fontsize=fontsize, fontweight=fontweight, color=labelcolor)
            j += 1
        i += 1
    return ax

def metapopviz(mp, loci='all', generation=-1, withlabels=True, fontsize=16, fontweight='bold', showaxes=False, figsize=(), offset=0.1, gcountpos='top', debug=False):
    if mp.stepping_stones:                          # test for stepping stones
        nX = nT = gridX = mp.number_of_populations
        nY = gridY = 1
        rowspan = 1
    elif mp.hybridzone:                                 # test for hybrid zone
        nX,nY = mp.hybridzoneshape                      # shape of the hybrid zone
        if nY%2 == 0:
            y_extension = 'even'
        else:
            y_extension = 'uneven'
        nHZ = nX*nY                     # number of populations in the hybrid zone
        nT = mp.number_of_populations   # total number of populations (hybrid zone + source populations)
        if not (nX*nY == nT-2):
            raise ValueError, "Invalid shape of metapopulation or total number of populations"
    
        if y_extension is 'even':
            gridY = 2*nY
            gridX = nX+2
            rowspan = 2
        else:
            gridY = nY
            gridX = nX+2
            rowspan = 1
    else:
        raise ShapeError, "Metapopulation must have a stepping stone or hybridzone shape"
        
    if debug:
        print "gridX:", gridX
        print "gridY:", gridY
        print "rowspan:",rowspan
    gridshape = (gridY, gridX)
    
    if not figsize:
        figsize = (3*gridX, 1.3*gridY)
    fig = plt.figure(figsize=figsize)
    if gcountpos == 'right':
        fig.text(0.99, 1.-1.5*offset, "generation", horizontalalignment='right', verticalalignment='top', fontsize=14, fontweight='bold', transform=fig.transFigure)
    elif isinstance(gcountpos, tuple):
        x,y = gcountpos
        fig.text(x+0.01,y, ".", horizontalalignment='right', verticalalignment='top', fontsize=6, fontweight='normal', color='0.99', transform=fig.transFigure)
        fig.text(x,y, "generation", horizontalalignment='right', verticalalignment='top', fontsize=14, fontweight='bold', transform=fig.transFigure)
    else:
        fig.text(1.-offset, 1.-offset, "generation", horizontalalignment='right', verticalalignment='top', fontsize=14, fontweight='bold', transform=fig.transFigure)
    ax = []
    
    if debug:
        print "positions on grid:"
    # left source:
    pos = gridY/2-1, 0
    if debug:
        print "\tleft source:", pos
    ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
    ax[-1].text(0.5,0.5, 'ls', transform=ax[-1].transAxes)
#    ax[-1].set_axis_off()
    for i in range(1,mp.number_of_populations-1):
        # pos = (y,x)
        pos = 2*((i-1)/nX), (i-1)%nX + 1
        if debug:
            print "\tpop%d: %s" % (i, str(pos))
        ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
        ax[-1].text(0.5,0.5, "pop. %d" % i, transform=ax[-1].transAxes)
#        ax[-1].set_axis_off()
    pos = gridY/2-1, gridX-1
    if debug:
        print "\tright source:", pos
    ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
    ax[-1].text(0.5,0.5, 'rs', transform=ax[-1].transAxes)
#    ax[-1].set_axis_off()
#    fig.subplots_adjust(bottom=offset,left=offset,right=1.-offset,top=1.-offset)
    
    g = generation                    
    for i in range(mp.number_of_populations):
        ax[i].cla()
        popviz(ax[i], mp.populations[i], loci=loci, generation=g, withlabels=withlabels, fontsize=fontsize, fontweight=fontweight, showaxes=showaxes)
    if g == -1:
        fig.texts[-1].set_text("%d" % mp.generation)
    else:
        fig.texts[-1].set_text("%d" % g)
    plt.draw()
    plt.show()
    return fig

def polar_to_rect(pcoords):
    """Convert polar coordinates into rectangular coordinates.
    pcoords = [theta,r]
    Return rcoords = [x,y]
    """
    theta,r = pcoords[0],pcoords[1]
    return [r*cos(theta), r*sin(theta)]

def shift(x,y, (dx,dy)):
    return [x+dx, y+dy]

def star(n=6, rin=0.25, rout=0.5, sh=(0,0), **kwargs):
    # ploar vertices:
    pverts = [[k*pi/n,k%2*rin+(k%2+1)*rout] for k in range(0,2*n)]
    rverts = map(polar_to_rect, pverts)
    shifted = [shift(rx,ry,sh) for rx,ry in rverts]
    return patches.Polygon(shifted, **kwargs)


def animate_axes_hybrid_zone(mp, loci='all', withlabels=True, withstar=False, starloc=[('pop1','P1')], figsize=(), offset=0.1, outfilename='animation', outfileformat='mp4', play=False, debug=False, **vizkwargs):
    """Depict all alleles."""
    try:
        global rcParams

        files = []
        
        sh = mp.hybridzoneshape
        nX,nY = sh                      # shape of the hybrid zone
        if nY%2 == 0:
            y_extension = 'even'
        else:
            y_extension = 'uneven'
        nHZ = nX*nY                     # number of populations in the hybrid zone
        nT = mp.number_of_populations   # total number of populations (hybrid zone + source populations)
        if not (nX*nY == nT-2):
            raise ValueError, "invalid shape of metapopulation or total number of populations"
            
        generations = list(mp.get_generations())
        
        if y_extension is 'even':
            gridY = 2*nY
            gridX = nX+2
            rowspan = 2
        else:
            gridY = nY
            gridX = nX+2
            rowspan = 1
        gridshape = (gridY, gridX)
        
        if not figsize:
            figsize = (3*gridX, 1.3*gridY)
        fig = plt.figure(figsize=figsize)
        fig.text(1.-offset, 1.-offset, "generation", horizontalalignment='right', verticalalignment='top', fontsize=14, fontweight='bold', transform=fig.transFigure)
        ax = []
        
        if debug:
            print "positions on grid:"
        # left source:
        pos = gridY/2-1, 0
        if debug:
            print "\tleft source:", pos
        ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
        ax[-1].text(0.5,0.5, 'ls', transform=ax[-1].transAxes)
    #    ax[-1].set_axis_off()
        for i in range(1,mp.number_of_populations-1):
            # pos = (y,x)
            pos = 2*((i-1)/nX), (i-1)%nX + 1
            if debug:
                print "\tpop%d: %s" % (i, str(pos))
            ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
            ax[-1].text(0.5,0.5, "pop. %d" % i, transform=ax[-1].transAxes)
    #        ax[-1].set_axis_off()
        pos = gridY/2-1, gridX-1
        if debug:
            print "\tright source:", pos
        ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
        ax[-1].text(0.5,0.5, 'rs', transform=ax[-1].transAxes)
        
        
        # find mutations:
        if withstar:
            stardict = {}
            for pop,loc in starloc:
                p = mp.__dict__[pop]
                gi = nonzero(p.history.get_sum(loc))[0][0]     # nonzero returns a tuple
                fadenumber = 8
                for j in range(fadenumber):
                    if not gi+j in stardict.keys():
                        stardict[gi+j] = [(mp.populations.index(p), loc, 1-j/float(fadenumber))]   # population index, star location (allele), alpha (transparency)
                    else:
                        stardict[gi+j].append( (mp.populations.index(p), loc, 1-j/float(fadenumber)) )
                
        for g in generations:
            for i in range(mp.number_of_populations):
                ws = False
                loc = ''
                alph = 1.
                ax[i].cla()
                gi = generations.index(g)
                if withstar and gi in stardict.keys():
                    for startuple in stardict[gi]:
                        if startuple[0] == i:
                            ws = True
                            popindex,loc,alph = startuple
                popviz(ax[i], mp.populations[i], loci=loci, generation=g, withlabels=withlabels, withstar=ws, starloc=loc, staralpha=alph, **vizkwargs)
            fig.texts[-1].set_text("%d" % g)
            plt.draw()
            plt.show()
            fname = '_tmp%05d.png' % gi
            print 'Saving frame', fname
            savefig(fname)
            files.append(fname)

        plt.close(fig)
        
        repeatlastframe = 20
        print 'Repeating last frame %d times:' % repeatlastframe
        lastfname = files[-1]
        for i in range( len(generations), len(generations)+repeatlastframe ):
            fname = "_tmp%05d.png" % i
            print fname
            shutil.copy(lastfname, fname)
            files.append(fname)
            
        print 'Trimming frames to remove whitespace:'
        cmd = "mogrify -type Grayscale -trim -bordercolor white -border 1 +repage _tmp*.png"
        print cmd
        mogrified = os.system(cmd)
        
        print "Making movie %s.%s - this make take a while" % (outfilename, outfileformat)
        cmd = "mencoder 'mf://_tmp*.png' -mf type=png:fps=10 -ovc lavc -lavcopts vcodec=%(codec)s -oac copy -o %(quote)s%(fname)s.%(format)s%(quote)s" % \
                {'fname': outfilename, \
                 'codec': moviecodecs[outfileformat], \
                 'format': outfileformat, \
                 'quote': '"'}
        print cmd
        mencoded = os.system(cmd)    # store return value in variable mencode

#    except:
#        # cleanup
#        if files:
#            for fname in files: os.remove(fname)
#        raise
        
    finally:
        # cleanup
        if files:
            print 'Saving first and last frame, cleaning up ...'
            os.rename(files[0], "%s_firstframe.png" % outfilename)
            os.rename(files[-1], "%s_lastframe.png" % outfilename)
            for fname in files[1:-1]:
                os.remove(fname)
        # play animation if mencoder was successful and play toggle is True
        if not mencoded and play:
            os.system("mplayer -fixed-vo -idle %(fname)s.%(format)s" % {'fname': outfilename, 'format': outfileformat})

def animate_stepping_stones(mp, loci='all', withlabels=True, withstar=False, starloc=[('pop1','P1')], figsize=(), offset=0.1, outfilename='animation', outfileformat='mp4', play=False, gcountpos='top', **vizkwargs):
    """Depict all alleles."""
    try:
        files = []
        
        if not mp.stepping_stones:
            raise ShapeError, "Metapopulation must have a stepping stone shape"
        nT = mp.number_of_populations   # total number of populations (hybrid zone + source populations)            
        generations = list(mp.get_generations())
        
        gridX = nT
        gridY = 1
        gridshape = (gridY,gridX)
        rowspan = 1
        
        if not figsize:
            figsize = (gridX, gridY)
        fig = plt.figure(figsize=figsize)
        if gcountpos == 'right':
            fig.text(0.99, 1.-offset, "generation", horizontalalignment='right', verticalalignment='top', transform=fig.transFigure, **vizkwargs)
        elif isinstance(gcountpos, tuple):
            x,y = gcountpos
            fig.text(x+0.01,y, ".", horizontalalignment='right', verticalalignment='top', fontsize=6, fontweight='normal', color='0.99', transform=fig.transFigure)  # to avoid shifting when imagemagick does its trimming
            fig.text(x,y, "generation", horizontalalignment='right', verticalalignment='top', transform=fig.transFigure, **vizkwargs)
        else:
            fig.text(1.-offset, 0.99, "generation", horizontalalignment='right', verticalalignment='top', transform=fig.transFigure, **vizkwargs)
        ax = []
        for i in range(nT):
            # pos = (y,x)
            pos = (0, i)
            ax.append( plt.subplot2grid(gridshape, pos, rowspan=rowspan) )
            ax[-1].text(0.5,0.5, "pop. %d" % i, transform=ax[-1].transAxes)
#            ax[-1].set_axis_off()
        fig.subplots_adjust(bottom=offset,left=offset,right=1.-offset,top=1.-offset)
        draw()
        
        # find mutations:
        if withstar:
            stardict = {}
            for pop,loc in starloc:
                p = mp.__dict__[pop]
                gi = nonzero(p.history.get_sum(loc))[0][0]     # nonzero returns a tuple
                fadenumber = 8
                for j in range(fadenumber):
                    if not gi+j in stardict.keys():
                        stardict[gi+j] = [(mp.populations.index(p), loc, 1-j/float(fadenumber))]   # population index, star location (allele), alpha (transparency)
                    else:
                        stardict[gi+j].append( (mp.populations.index(p), loc, 1-j/float(fadenumber)) )
        
        repeatfirstframe = 10
        repeatlastframe = 20
        lastgi = len(generations)-1
        print 'Saving frames:'
        for g in generations:
            gi = generations.index(g)
            for i in range(nT):
                ws = False
                loc = ''
                alph = 1.
                ax[i].cla()
                if withstar and gi in stardict.keys():
                    for startuple in stardict[gi]:
                        if startuple[0] == i:
                            ws = True
                            popindex,loc,alph = startuple
                popviz(ax[i], mp.populations[i], loci=loci, generation=g, withlabels=withlabels, withstar=ws, starloc=loc, staralpha=alph, **vizkwargs)
            fig.texts[-1].set_text("%d" % g)
            plt.draw()
            plt.show()
            if gi == 0:
                for j in range(repeatfirstframe):
                    fname = '_tmp%05d.png' % j
                    print fname
                    savefig(fname)
                    files.append(fname)
            fname = '_tmp%05d.png' % (gi+repeatfirstframe)    # use the index here
            print fname
            savefig(fname)
            files.append(fname)
            if gi == lastgi:
                for j in range(1,repeatlastframe+1):
                    fname = '_tmp%05d.png' % (gi+repeatfirstframe+j)
                    print fname
                    savefig(fname)
                    files.append(fname)

        plt.close(fig)
        
#        repeatfirstframe = 10
#        print 'Repeating first frame %d times:' % repeatfirstframe
#        firstfname = files[0]
#        for i in range( len(generations), len(generations)+repeatlastframe ):
#            fname = "_tmp%05d.png" % i
#            print fname
#            shutil.copy(lastfname, fname)
#            files.append(fname)

#        repeatlastframe = 20
#        print 'Repeating last frame %d times:' % repeatlastframe
#        lastfname = files[-1]
#        for i in range( len(generations), len(generations)+repeatlastframe ):
#            fname = "_tmp%05d.png" % i
#            print fname
#            shutil.copy(lastfname, fname)
#            files.append(fname)
            
            
        print 'Trimming frames to remove whitespace:'
        cmd = "mogrify -type Grayscale -trim -bordercolor white -border 1 +repage _tmp*.png"
        print cmd
        mogrified = os.system(cmd)
        
        print "Making movie %s.%s - this may take a while" % (outfilename, outfileformat)
#        artist = 'Matthias Flor'
#        title = "Stable Wolbachia infection polymorphism"
#        comment = "Invasive wave front of Wolbachia stopped by weak migration between two populations"
        cmd = "mencoder 'mf://_tmp*.png' -mf type=png:fps=10 -ovc lavc -lavcopts vcodec=%(codec)s -oac copy -o %(quote)s%(fname)s.%(format)s%(quote)s" % \
                {'fname': outfilename, \
                 'codec': moviecodecs[outfileformat], \
                 'format': outfileformat, \
                 'quote': '"'}
        print cmd
        mencoded = os.system(cmd)    # store return value in variable mencode
#        print 'mencode:', mencode

#    except:
#        # cleanup
#        if files:
#            for fname in files: os.remove(fname)
#        raise
        
    finally:
        # cleanup
        if files:
            print files
            print 'Saving first and last frame, cleaning up ...'
#            os.rename(files[0], "%s_firstframe.png" % outfilename)
#            os.rename(files[-1], "%s_lastframe.png" % outfilename)
#            for fname in files[1:-1]:
#                os.remove(fname)
        # play animation if mencoder was successful and play toggle is True
        if not mencoded and play:
            os.system("mplayer -fixed-vo -idle %(fname)s.%(format)s" % {'fname': outfilename, 'format': outfileformat})
        
if __name__ == '__main__':
    main()        
        
